import os
import torch
from torch.utils.data import Dataset
from PIL import Image
import numpy as np

class MyDataset(Dataset):
    def __init__(self,path):
        self.path = path
        self.dataset = []
        self.dataset.extend(os.listdir(path))

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self,index):
        label = torch.Tensor(np.array([int(self.dataset[index][0])]))
        img_path = os.path.join(self.path,self.dataset[index])
        img = Image.open(img_path)
        img_data = torch.Tensor(np.array(img) / 255 - 0.5)
        return img_data,label

if  __name__ == '__main__':
    ImgTakeOut = ("img")
    x = ImgTakeOut[9][0].numpy()
    y = ImgTakeOut[9][0].numpy()

    img_data = np.array((x+0.5)*255,dtype=np.int8)
    img = Image.fromarray(img_data,"RGB")
    img.show()
  



